package cn.xiaopengstack.infrastructure.adapter.repository;

import cn.xiaopengstack.domain.agent.adapter.repository.IAgentRepository;
import cn.xiaopengstack.domain.agent.model.valobj.*;
import cn.xiaopengstack.infrastructure.dao.*;
import cn.xiaopengstack.infrastructure.dao.po.*;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.JSONValidator;
import com.alibaba.fastjson.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Repository;
import org.springframework.util.CollectionUtils;

import javax.validation.constraints.NotNull;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

import static cn.xiaopengstack.domain.agent.model.valobj.AiAgentEnumVO.*;

/**
 * @author jiangyangang
 */
@Slf4j
@Repository
public class AgentRepository implements IAgentRepository {
    @Resource
    private AiClientMapper aiClientMapper;
    @Resource
    private AiClientApiMapper aiClientApiMapper;
    @Resource
    private AiClientModelMapper aiClientModelMapper;
    @Resource
    private AiClientSystemPromptMapper aiClientSystemPromptMapper;
    @Resource
    private AiClientToolMcpMapper aiClientToolMcpMapper;
    @Resource
    private AiClientConfigMapper aiClientConfigMapper;
    @Resource
    private AiClientAdvisorMapper aiClientAdvisorMapper;
    @Resource
    private AiAgentFlowConfigMapper aiAgentFlowConfigMapper;
    @Resource
    private AiAgentTaskScheduleMapper aiAgentTaskScheduleMapper;
    @Resource
    private AiAgentMapper aiAgentMapper;

    @Override
    public List<AiClientApiVO> queryAiClientApiVOListByClientIds(List<String> clients) {
        if (CollectionUtils.isEmpty(clients)) {
            return List.of();
        }

        ArrayList<AiClientApiVO> result = new ArrayList<>();
        // 由于clientId是配置在更高层级的配置，所以需要从上往下找。
        for (String clientId : clients) {
            List<AiClientConfig> aiClientConfigs = aiClientConfigMapper.queryBySourceAndTarget(AI_CLIENT.getCode(), clientId, AI_CLIENT_MODEL.getCode());

            // 获取对话模型
            for (AiClientConfig aiClientConfig : aiClientConfigs) {

                String modelId = aiClientConfig.getTargetId();
                AiClientModel aiClientModel = aiClientModelMapper.queryByModelId(modelId);
                if (Objects.isNull(aiClientModel)) {
                    log.warn("基础数据仓储查询，agent client【{}】下没有对应的对话模型，跳过", clientId);
                    continue;
                }
                String apiId = aiClientModel.getApiId();
                // 获取对话模型需要的api信息
                AiClientApi aiClientApi = aiClientApiMapper.queryByApiId(apiId);
                if (Objects.isNull(aiClientApi)) {
                    log.warn("基础数据仓储查询，agent client【{}】的对话模型【{}-{}】没有对应的API配置，跳过", clientId, aiClientModel.getModelId(), aiClientModel.getModelName());
                    continue;
                }

                AiClientApiVO apiVo = AiClientApiVO.builder()
                        .apiKey(aiClientApi.getApiKey())
                        .baseUrl(aiClientApi.getBaseUrl())
                        .completionsPath(aiClientApi.getCompletionsPath())
                        .embeddingsPath(aiClientApi.getEmbeddingsPath())
                        .apiId(aiClientApi.getApiId())
                        .build();
                if (result.stream().noneMatch(vo -> vo.getApiId().equals(apiVo.getApiId()))) {
                    result.add(apiVo);
                }
            }
        }

        return result;
    }

    @Override
    public List<AiClientModelVO> aiClientModelVOByClientIds(List<String> clients) {
        if (CollectionUtils.isEmpty(clients)) {
            return List.of();
        }

        ArrayList<AiClientModelVO> result = new ArrayList<>();
        // 由于clientId是配置在更高层级的配置，所以需要从上往下找。
        for (String clientId : clients) {
            List<AiClientConfig> aiClientConfigs = aiClientConfigMapper.queryBySourceAndTarget(AI_CLIENT.getCode(), clientId, AI_CLIENT_MODEL.getCode());

            // 获取对话模型
            for (AiClientConfig aiClientConfig : aiClientConfigs) {

                String modelId = aiClientConfig.getTargetId();
                AiClientModel aiClientModel = aiClientModelMapper.queryByModelId(modelId);
                if (Objects.isNull(aiClientModel)) {
                    log.warn("基础数据仓储查询，agent client【{}】下没有对应的对话模型，跳过", clientId);
                    continue;
                }

                List<AiClientConfig> mcpList = aiClientConfigMapper.queryBySourceAndTarget(AI_CLIENT_MODEL.getCode(), modelId, AI_CLIENT_TOOL_MCP.getCode());
                List<String> mcpIds = List.of();
                if (!CollectionUtils.isEmpty(mcpList)) {
                    mcpIds = mcpList.stream().map(AiClientConfig::getTargetId).toList();
                }

                AiClientModelVO modelVO = AiClientModelVO.builder()
                        .modelName(aiClientModel.getModelName())
                        .apiId(aiClientModel.getApiId())
                        .modelType(aiClientModel.getModelType())
                        .modelId(aiClientModel.getModelId())
                        .toolMcpIds(mcpIds)
                        .build();
                if (result.stream().noneMatch(vo -> vo.getModelId().equals(modelVO.getModelId()))) {
                    result.add(modelVO);
                }
            }
        }
        return result;
    }

    @Override
    public List<AiClientToolMcpVO> aiClientToolMcpVOByClientIds(List<String> clients) {
        if (CollectionUtils.isEmpty(clients)) {
            return List.of();
        }

        ArrayList<AiClientToolMcpVO> result = new ArrayList<>();
        // 由于clientId是配置在更高层级的配置，所以需要从上往下找。
        for (String clientId : clients) {
            List<AiClientConfig> aiClientConfigs = aiClientConfigMapper.queryBySourceAndTarget(AI_CLIENT.getCode(), clientId, AI_CLIENT_TOOL_MCP.getCode());

            // 获取对话模型
            for (AiClientConfig aiClientConfig : aiClientConfigs) {

                String mcpId = aiClientConfig.getTargetId();

                AiClientToolMcp aiClientToolMcp = aiClientToolMcpMapper.queryByMcpId(mcpId);
                if (Objects.isNull(aiClientToolMcp)) {
                    log.warn("基础数据仓储查询，agent mcp【{}】下没有对应的Mcp配置，跳过", clientId);
                    continue;
                }

                AiClientToolMcpVO mcpVO = AiClientToolMcpVO.builder()
                        .mcpName(aiClientToolMcp.getMcpName())
                        .requestTimeout(aiClientToolMcp.getRequestTimeout())
                        .mcpId(aiClientToolMcp.getMcpId())
                        .transportConfig(aiClientToolMcp.getTransportConfig())
                        .transportType(aiClientToolMcp.getTransportType())
                        .build();

                try {
                    if (McpTransportTypeEnumVO.sse.getType().equals(aiClientToolMcp.getTransportType())) {
                        ObjectMapper objectMapper = new ObjectMapper();
                        AiClientToolMcpVO.TransportConfigSse transportConfigSse = objectMapper.readValue(aiClientToolMcp.getTransportConfig(), AiClientToolMcpVO.TransportConfigSse.class);
                        mcpVO.setTransportConfigSse(transportConfigSse);
                    }
                    if (McpTransportTypeEnumVO.stdio.getType().equals(aiClientToolMcp.getTransportType())) {
                        Map<String, AiClientToolMcpVO.TransportConfigStdio.Stdio> stdioMap = JSONObject.parseObject(aiClientToolMcp.getTransportConfig(), new TypeReference<>(){});
                        AiClientToolMcpVO.TransportConfigStdio transportConfigStdio = new AiClientToolMcpVO.TransportConfigStdio();
                        transportConfigStdio.setStdio(stdioMap);
                        mcpVO.setTransportConfigStdio(transportConfigStdio);
                    }
                } catch (Exception e) {
                    log.error("解析传输配置失败: {}", e.getMessage(), e);
                }

                if (result.stream().noneMatch(vo -> vo.getMcpId().equals(mcpVO.getMcpId()))) {
                    result.add(mcpVO);
                }
            }
        }
        return result;
    }

    @Override
    public List<AiClientSystemPromptVO> aiClientSystemPromptVOByClientIds(List<String> clients) {
        if (CollectionUtils.isEmpty(clients)) {
            return List.of();
        }

        ArrayList<AiClientSystemPromptVO> result = new ArrayList<>();
        // 由于clientId是配置在更高层级的配置，所以需要从上往下找。
        for (String clientId : clients) {
            List<AiClientConfig> aiClientConfigs = aiClientConfigMapper.queryBySourceAndTarget(AI_CLIENT.getCode(), clientId, AI_CLIENT_SYSTEM_PROMPT.getCode());

            // 获取对话模型
            for (AiClientConfig aiClientConfig : aiClientConfigs) {
                String promptId = aiClientConfig.getTargetId();

                AiClientSystemPrompt aiClientSystemPrompt = aiClientSystemPromptMapper.queryByPromptId(promptId);
                if (Objects.isNull(aiClientSystemPrompt)) {
                    log.warn("基础数据仓储查询，agent client【{}】下没有对应的提示词配置，跳过", clientId);
                    continue;
                }

                AiClientSystemPromptVO promptVO = AiClientSystemPromptVO.builder()
                        .promptContent(aiClientSystemPrompt.getPromptContent())
                        .promptName(aiClientSystemPrompt.getPromptName())
                        .description(aiClientSystemPrompt.getDescription())
                        .promptId(aiClientSystemPrompt.getPromptId())
                        .build();

                if (result.stream().noneMatch(vo -> vo.getPromptId().equals(promptVO.getPromptId()))) {
                    result.add(promptVO);
                }
            }
        }
        return result;
    }

    @Override
    public Map<String, AiClientSystemPromptVO> queryAiClientSystemPromptMapByClientIds(List<String> clients) {

        if (CollectionUtils.isEmpty(clients)) {
            return Map.of();
        }

        Map<String, AiClientSystemPromptVO> result = new HashMap<>();
        // 由于clientId是配置在更高层级的配置，所以需要从上往下找。
        for (String clientId : clients) {
            List<AiClientConfig> aiClientConfigs = aiClientConfigMapper.queryBySourceAndTarget(AI_CLIENT.getCode(), clientId, AI_CLIENT_SYSTEM_PROMPT.getCode());

            // 获取对话模型
            for (AiClientConfig aiClientConfig : aiClientConfigs) {
                String promptId = aiClientConfig.getTargetId();

                AiClientSystemPrompt aiClientSystemPrompt = aiClientSystemPromptMapper.queryByPromptId(promptId);
                if (Objects.isNull(aiClientSystemPrompt)) {
                    log.warn("基础数据仓储查询，agent client【{}】下没有对应的提示词配置，跳过", clientId);
                    continue;
                }

                AiClientSystemPromptVO promptVO = AiClientSystemPromptVO.builder()
                        .promptContent(aiClientSystemPrompt.getPromptContent())
                        .promptName(aiClientSystemPrompt.getPromptName())
                        .description(aiClientSystemPrompt.getDescription())
                        .promptId(aiClientSystemPrompt.getPromptId())
                        .build();

                result.put(promptId, promptVO);
            }
        }
        return result;
    }

    @Override
    public List<AiClientAdvisorVO> aiClientAdvisorVOByClientIds(List<String> clients) {
        if (CollectionUtils.isEmpty(clients)) {
            return List.of();
        }

        ArrayList<AiClientAdvisorVO> result = new ArrayList<>();
        // 由于clientId是配置在更高层级的配置，所以需要从上往下找。
        for (String clientId : clients) {
            List<AiClientConfig> aiClientConfigs = aiClientConfigMapper.queryBySourceAndTarget(AI_CLIENT.getCode(), clientId, AI_CLIENT_ADVISOR.getCode());

            // 获取对话模型
            for (AiClientConfig aiClientConfig : aiClientConfigs) {
                String advisorId = aiClientConfig.getTargetId();

                AiClientAdvisor aiClientAdvisor = aiClientAdvisorMapper.queryByAdvisorId(advisorId);
                if (Objects.isNull(aiClientAdvisor)) {
                    log.warn("基础数据仓储查询，agent client【{}】下没有对应的助手配置，跳过", clientId);
                    continue;
                }

                // 3. 解析extParam中的配置
                AiClientAdvisorVO.ChatMemory chatMemory = null;
                AiClientAdvisorVO.RagAnswer ragAnswer = null;
                try {
                    if (JSONValidator.from(aiClientAdvisor.getExtParam()).validate()) {
                        if (AiAdvisorTypeEnumVO.ChatMemory.getType().equals(aiClientAdvisor.getAdvisorType())){
                            chatMemory = JSONObject.parseObject(aiClientAdvisor.getExtParam(), AiClientAdvisorVO.ChatMemory.class);
                        }
                        if (AiAdvisorTypeEnumVO.RagAnswer.getType().equals(aiClientAdvisor.getAdvisorType())) {
                            ragAnswer = JSONObject.parseObject(aiClientAdvisor.getExtParam(), AiClientAdvisorVO.RagAnswer.class);
                        }
                    }
                } catch (Exception e) {
                    log.warn("基础数据仓储查询, 装载顾问JSON配置异常, advisorId={}", aiClientAdvisor.getAdvisorId());
                }

                AiClientAdvisorVO advisorVO = AiClientAdvisorVO.builder()
                        .advisorId(aiClientAdvisor.getAdvisorId())
                        .advisorName(aiClientAdvisor.getAdvisorName())
                        .advisorType(aiClientAdvisor.getAdvisorType())
                        .orderNum(aiClientAdvisor.getOrderNum())
                        .chatMemory(chatMemory)
                        .ragAnswer(ragAnswer)
                        .build();


                if (result.stream().noneMatch(vo -> vo.getAdvisorId().equals(advisorVO.getAdvisorId()))) {
                    result.add(advisorVO);
                }
            }
        }
        return result;
    }

    @Override
    public List<AiClientVO> aiClientVOByClientIds(List<String> clients) {
        if (CollectionUtils.isEmpty(clients)) {
            return List.of();
        }
        ArrayList<AiClientVO> result = new ArrayList<>();

        for (String clientId : clients) {
            List<AiClientConfig> aiClientConfigs = aiClientConfigMapper.queryBySourceTypeAndId(AI_CLIENT.getCode(), clientId);
            List<String> promptIdList = new ArrayList<>();
            List<String> mcpIdList = new ArrayList<>();
            List<String> advisorIdList = new ArrayList<>();
            String modelId = null;
            AiClient aiClient = aiClientMapper.queryByClientId(clientId);

            for (AiClientConfig aiClientConfig : aiClientConfigs) {
                AiClientConfigTargetTypeEnumVO typeEnumVO = AiClientConfigTargetTypeEnumVO.valueOf(aiClientConfig.getTargetType());
                switch (typeEnumVO) {
                    case prompt:
                        promptIdList.add(aiClientConfig.getTargetId());
                        break;
                    case tool_mcp:
                        mcpIdList.add(aiClientConfig.getTargetId());
                        break;
                    case model:
                        modelId = aiClientConfig.getTargetId();
                        break;
                    case advisor:
                        advisorIdList.add(aiClientConfig.getTargetId());
                        break;
                }
            }

            AiClientVO clientVO = AiClientVO.builder()
                    .clientName(aiClient.getClientName())
                    .clientId(aiClient.getClientId())
                    .description(aiClient.getDescription())
                    .modelId(modelId)
                    .advisorIdList(advisorIdList)
                    .promptIdList(promptIdList)
                    .mcpIdList(mcpIdList)
                    .build();
            if (result.stream().noneMatch(vo -> vo.getClientId().equals(clientVO.getClientId()))) {
                result.add(clientVO);
            }
        }

        return result;
    }

    @Override
    public List<AiClientApiVO> queryAiClientApiVOListByModelIds(List<String> modelIdList) {
        if (CollectionUtils.isEmpty(modelIdList)) {
            return List.of();
        }

        ArrayList<AiClientApiVO> result = new ArrayList<>();
        // 由于clientId是配置在更高层级的配置，所以需要从上往下找。
        for (String modelId : modelIdList) {
            AiClientModel aiClientModel = aiClientModelMapper.queryByModelId(modelId);
            if (Objects.isNull(aiClientModel)) {
                log.warn("基础数据仓储查询，agent modelId【{}】下没有对应的对话模型，跳过", modelId);
                continue;
            }
            String apiId = aiClientModel.getApiId();
            // 获取对话模型需要的api信息
            AiClientApi aiClientApi = aiClientApiMapper.queryByApiId(apiId);
            if (Objects.isNull(aiClientApi)) {
                log.warn("基础数据仓储查询，对话模型【{}-{}】没有对应的API配置，跳过", aiClientModel.getModelId(), aiClientModel.getModelName());
                continue;
            }

            AiClientApiVO apiVo = AiClientApiVO.builder()
                    .apiKey(aiClientApi.getApiKey())
                    .baseUrl(aiClientApi.getBaseUrl())
                    .completionsPath(aiClientApi.getCompletionsPath())
                    .embeddingsPath(aiClientApi.getEmbeddingsPath())
                    .apiId(aiClientApi.getApiId())
                    .build();
            if (result.stream().noneMatch(vo -> vo.getApiId().equals(apiVo.getApiId()))) {
                result.add(apiVo);
            }
        }

        return result;
    }

    @Override
    public List<AiClientModelVO> AiClientModelVOByModelIds(List<String> modelIdList) {
        if (CollectionUtils.isEmpty(modelIdList)) {
            return List.of();
        }

        ArrayList<AiClientModelVO> result = new ArrayList<>();
        // 由于clientId是配置在更高层级的配置，所以需要从上往下找。
        for (String modelId : modelIdList) {
            AiClientModel aiClientModel = aiClientModelMapper.queryByModelId(modelId);
            if (Objects.isNull(aiClientModel)) {
                log.warn("基础数据仓储查询，agent model【{}】下没有对应的对话模型，跳过", modelId);
                continue;
            }
            AiClientModelVO modelVO = AiClientModelVO.builder()
                    .modelName(aiClientModel.getModelName())
                    .apiId(aiClientModel.getApiId())
                    .modelType(aiClientModel.getModelType())
                    .modelId(aiClientModel.getModelId())
                    .build();
            if (result.stream().noneMatch(vo -> vo.getApiId().equals(modelVO.getApiId()))) {
                result.add(modelVO);
            }
        }
        return result;
    }

    @Override
    public Map<String, AiAgentClientFlowConfigVO> queryAiAgentClientFlowConfig(@NotNull String aiAgentId) {

        List<AiAgentFlowConfig> agentFlowConfigs = aiAgentFlowConfigMapper.queryByAgentId(aiAgentId);

        // 转换为Map结构，key为clientId，value为AiAgentClientFlowConfigVO
        Map<String, AiAgentClientFlowConfigVO> result = new HashMap<>();
        for (AiAgentFlowConfig agentFlowConfig : agentFlowConfigs) {

            AiAgentClientFlowConfigVO clientFlowConfigVO = AiAgentClientFlowConfigVO.builder()
                    .clientType(agentFlowConfig.getClientType())
                    .sequence(agentFlowConfig.getSequence())
                    .clientName(agentFlowConfig.getClientName())
                    .stepPrompt(agentFlowConfig.getStepPrompt())
                    .clientId(agentFlowConfig.getClientId())
                    .build();

            result.put(agentFlowConfig.getClientType(), clientFlowConfigVO);
        }

        return result;
    }

    @Override
    public AiAgentVo queryAgentV0ById(String agentId) {
        AiAgent agent = aiAgentMapper.queryByAgentId(agentId);
        assert null != agent;

        AiAgentVo aiAgentVo = AiAgentVo.builder()
                .agentId(agent.getAgentId())
                .agentName(agent.getAgentName())
                .channel(agent.getChannel())
                .strategy(agent.getStrategy())
                .description(agent.getDescription())
                .build();

        return aiAgentVo;
    }
}
